# -*- coding: utf-8 -*-
#Created on Mon Aug 6 9:49:15 2018
#@author: JockJo

from pymongo import MongoClient


class Mongo():
    collection_name = ''
    mongo_db = ''
    client=MongoClient()

    #@client: 指定client，ip地址
    #@port：指定端口
    #@collection_name:如果只对一个collection做操作，可以初始化的时候指定collection
    #mongo_db:指定数据库，不能为空
    def __init__(self, client=None, port=None, collection_name=None, mongo_db=None):
        if mongo_db is None:
            print('[mongodb error] mongodb is not specific')
            return

        self.collection_name = collection_name
        self.mongo_db = mongo_db

        if client is None:
            self.client = MongoClient('localhost', 27017)     # default connection
        else:
            self.client = MongoClient(client, port)            # specific connection
        self.db = self.client[self.mongo_db]

    def store_item(self, collection_name, item):
        collection_name = collection_name if collection_name is not None else self.collection_name
        if collection_name is None:
            print('[mongodb error] collection is not specific')
        record = {'applicants':item['us_applicants'],'inventors':item['us_inventors'],
                  'agents':item['us_agents'],'assignees':item['us_assignee'],
                  'abstract':item['abstract'],'description':item['description'],'claims':item['claim'],
                  'title':item['title'],'base_info':item['base_info'],'application_reference':item['us_application_reference'],
                  'classification':item['us_classification']}
        self.db[collection_name].update({'_id': item['doc_number']},{'$set': record},upsert=True)

    def get_item(self, collection_name, item, field):
        record = self.db[self.collection_name].find_one({'field': item})
        if record:
            return not None
        else:
            return None

    def get_company(self, collection_name):
        company = self.db[collection_name].find()
        company_list = []
        for c in company:
            if c['Conme'] is not '' and c['Stknme'] is not '':
                if type(c['Stkcd']) is type(''):
                    query_symbol = 'sh'+c['Stkcd'] if c['Stkcd'][0] == '6' and c['Stkcd'][1] == '0' else 'sz'+c['Stkcd']
                else:
                    query_symbol = 'sz'+str(c['Stkcd']) if c['Stkcd'] < 600000 else 'sh'+str(c['Stkcd'])
                company_list.append([c['Stknme'], c['Conme'], c['Stkcd'],query_symbol])  # 证券名称,公司全称,证券代码
        return company_list

    def close(self):
        self.client.close()

